import autolens as al
import autolens.plot as aplt
import autofit as af
import numpy as np
from matplotlib import pyplot as plt
from astropy.io import fits
from os import path

'''necessary parameters'''
###source and lens redshift
redshift_l = 1.7
redshift_s = 2.5579

###the position of the four point sources, to accelerate the fitting
positions = np.array([[-0.327,-0.0057*15],[-0.152,-0.0550*15],[0.723,-0.0295*15],[0.398,0.0283*15]])-np.array([[0.244,-0.0172*15]])
positions = al.Grid2DIrregular(positions)

###the folder that output the results
dataname = "340GHz"
testid = "cl_parametric"

'''run the fitting'''
###input the data image
imaging = al.Imaging.from_fits(
    image_path=f"./data/340GHz_cut_f8.fits",
    noise_map_path=f"./data/340GHz_cut_n8.fits",
    psf_path=f"./data/340GHz_psf4.fits",
    pixel_scales=0.03,
)

###add a mask
mask = al.Mask2D.circular(
    shape_native=imaging.shape_native, pixel_scales=imaging.pixel_scales, radius=1.1
)
imaging = imaging.apply_mask(mask=mask)


'''the model and the fitting'''
###lens and source model
lens_galaxy_model = af.Model(al.Galaxy, redshift=redshift_l, mass=al.mp.EllPowerLaw, shear=al.mp.ExternalShear)
source_galaxy_model = af.Model(al.Galaxy, redshift=redshift_s, bulge=al.lp.EllSersic)

###fix the centre of the lens
lens_galaxy_model.mass.centre_0 = 0.001
lens_galaxy_model.mass.centre_1 = -0.003

model = af.Collection(
    galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)
)

###set a threshold
settings_lens = al.SettingsLens(positions_threshold=0.1)

###mcmc fitting
search = af.DynestyStatic(
    path_prefix=path.join("howtolens"),
    name=testid,
    unique_tag=dataname,
    nlive=200,
    number_of_cores=8,
    walks=10 
)

analysis = al.AnalysisImaging(dataset=imaging, positions=positions, settings_lens=settings_lens)

result = search.fit(model=model, analysis=analysis)

